"""
控制模式实现
"""

import re
from rich.console import Console
from rich.prompt import Prompt
from rich.panel import Panel
from relay_api import RelayAPI
from ssh_manager import SSHManager

console = Console()

class ControlMode:
    def __init__(self):
        self.api = RelayAPI()
        self.ssh_manager = SSHManager()
    
    def start(self):
        """启动控制模式"""
        console.print("\n[bold blue]启动控制模式...[/bold blue]")
        
        # 获取密钥
        key = self._get_key_from_user()
        if not key:
            return False
        
        # 验证密钥并获取连接信息
        console.print("正在验证密钥...")
        result = self.api.connect(key)
        
        if result.get("status") != "success":
            console.print(f"[bold red]连接失败: {result.get('message', '未知错误')}[/bold red]")
            return False
        
        # 显示目标信息
        self._display_target_info(result)
        
        # 连接到目标
        relay_host = result["relay_host"]
        tunnel_port = result["tunnel_port"]
        
        console.print(f"正在连接到目标主机...")
        console.print(f"[dim]通过 {relay_host}:{tunnel_port} 建立连接[/dim]")
        console.print("\n[bold green]即将进入SSH会话，输入 'exit' 退出连接[/bold green]\n")
        
        # 建立SSH连接
        success = self.ssh_manager.connect_to_target(relay_host, tunnel_port)
        
        if success:
            console.print("\n[bold green]连接已断开[/bold green]")
        else:
            console.print("\n[bold red]连接失败[/bold red]")
        
        return success
    
    def _get_key_from_user(self) -> str:
        """从用户获取密钥"""
        while True:
            key = Prompt.ask(
                "[bold cyan]请输入8位连接密钥[/bold cyan]",
                default="",
                show_default=False
            ).strip().upper()
            
            if not key:
                console.print("[bold yellow]已取消连接[/bold yellow]")
                return ""
            
            # 验证密钥格式
            if self._validate_key_format(key):
                return key
            else:
                console.print("[bold red]密钥格式错误！请输入8位数字字母组合[/bold red]")
    
    def _validate_key_format(self, key: str) -> bool:
        """验证密钥格式"""
        if len(key) != 8:
            return False
        
        # 检查是否只包含数字和大写字母
        if not re.match(r'^[A-Z0-9]{8}$', key):
            return False
        
        # 检查是否至少包含一个数字和一个字母
        has_digit = any(c.isdigit() for c in key)
        has_alpha = any(c.isalpha() for c in key)
        
        return has_digit and has_alpha
    
    def _display_target_info(self, connection_info: dict):
        """显示目标信息"""
        client_info = connection_info.get("client_info", "未知")
        relay_host = connection_info.get("relay_host", "未知")
        tunnel_port = connection_info.get("tunnel_port", "未知")
        ssh_port = connection_info.get("ssh_port", 22)
        
        info_text = f"""
[bold white]目标主机信息:[/bold white]

[bold cyan]系统信息:[/bold cyan] {client_info}
[bold cyan]中继服务器:[/bold cyan] {relay_host}
[bold cyan]隧道端口:[/bold cyan] {tunnel_port}
[bold cyan]SSH端口:[/bold cyan] {ssh_port}

[dim]连接将通过中继服务器建立安全隧道[/dim]
        """.strip()
        
        panel = Panel(
            info_text,
            title="[bold blue]连接信息[/bold blue]",
            border_style="blue",
            padding=(1, 2)
        )
        
        console.print(panel)

